import os
import json
from dashscope import Assistants, Messages, Runs, Threads, Application
from prompt import *


def parse_streaming_response(response):
    complete_response = ""
    for chunk in response:
        complete_response += chunk
    return complete_response


def init_thread():
    new_thread = Threads.create()
    return {"thread_id": new_thread.id}


router_agent = Assistants.create(
    model="qwen-plus",
    name='引导员，路由器',
    description='你是一个商城的引导员，负责将用户问题路由到不同的导购员。',
    instructions=ROUTER_AGENT_INSTRUCTION
)

mobilephone_guide_agent = Assistants.create(
    model="qwen-max",
    name='手机导购',
    description='你是一个手机导购，你需要询问顾客想要什么参数的手机。',
    instructions=MOBILEPHONE_GUIDE_AGENT_INSTRUCTION
)

fridge_guide_agent = Assistants.create(
    model="qwen-max",
    name='冰箱导购',
    description='你是一个冰箱导购，你需要询问顾客想要什么参数的冰箱。',
    instructions=FRIDGE_GUIDE_AGENT_INSTRUCTION
)

tv_guide_agent = Assistants.create(
    model="qwen-max",
    name='电视导购',
    description='你是一个电视导购，你需要询问顾客想要什么参数的电视。',
    instructions=TV_GUIDE_AGENT_INSTRUCTION
)

# 定义一个映射表，将agent的名称映射到agent的id
agent_map = {
    "意图分类": router_agent.id,
    "手机": mobilephone_guide_agent.id,
    "冰箱": fridge_guide_agent.id,
    "电视机": tv_guide_agent.id
}


def get_agent_response(agent_name, input_prompt, thread_id):
    agent_id = agent_map[agent_name]
    message = Messages.create(thread_id=thread_id, content=input_prompt)
    run = Runs.create(thread_id=thread_id, assistant_id=agent_id, stream=True)
    for event, msg in run:
        try:
            content = msg['delta']['content']['text']['value']
            yield content
        except:
            pass



def get_rag_response(input_prompt):
    response = Application.call(
        app_id=os.environ['BAILIAN_APP_ID'],
        prompt=input_prompt,
        stream=True,
        incremental_output=True
    )
    for chunk in response:
        yield chunk.output.text



def chat(input_prompt, thread_id):

    router_agent_response = get_agent_response(agent_name="意图分类", input_prompt=input_prompt, thread_id=thread_id)
    classification_result = parse_streaming_response(router_agent_response)
    print('result:', classification_result)

    response_json = {
        "content": "",
    }


    if classification_result == "其他":
        response_json["content"] = "不好意思，我没有理解您的问题，能换个表述方式么？"
        response_json['current_agent'] = classification_result
        response_json['thread_id'] = thread_id
        yield f"{json.dumps(response_json)}\n\n"
 
    else:
        response = get_agent_response(agent_name=classification_result, input_prompt=input_prompt, thread_id=thread_id)
        whole_response = ""
        for chunk in response:
            whole_response += chunk
            response_json["content"] = chunk
            response_json['current_agent'] = classification_result
            response_json['thread_id'] = thread_id
            yield f"{json.dumps(response_json)}\n\n"
